#!/usr/bin/env python3
"""
轻量级 Modbus TCP 测试服务器

特性：
  * 支持功能码 0x03（Read Holding Registers）与 0x04（Read Input Registers）
  * 默认提供从地址 0 起连续 N 个寄存器（默认 200 个），值为索引 + 填充值
  * 提供简单的自增模拟：每次请求后寄存器整体 +1

示例：
    python tools/modbus_tcp_sim.py --host 0.0.0.0 --port 1502 --count 300 --fill 100

    # 让寄存器值持续递增（模拟实时数据）
    python tools/modbus_tcp_sim.py --auto-increment
"""

import argparse
import logging
import socket
import socketserver
import struct
import math
import threading
import time
from typing import Dict, List, Optional

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)


def make_zkzf_update() -> callable:
    """创建一个动态更新函数，模拟中科智赋的关键寄存器波动。"""

    last_time = time.time()
    daily_energy = 5000  # 0.1 kWh 单位
    yearly_energy = 28000
    warehouse_base = 260

    def updater(store: "ModbusDataStore") -> None:
        nonlocal last_time, daily_energy, yearly_energy, warehouse_base

        now = time.time()
        dt = now - last_time
        if dt <= 0:
            dt = store.interval
        last_time = now

        # 模拟光伏功率（单位 0.1kW）
        pv_kw = 120 + 45 * math.sin(now / 18.0)
        pv_reg = max(0, int(pv_kw * 10))
        store.set_value(35, pv_reg)
        store.set_value(40, pv_reg + 150)  # 直流母线功率略高

        # 直流母线电压、电流
        store.set_value(38, 750 + int(20 * math.sin(now / 22.0)))
        store.set_value(39, 150 + int(30 * math.cos(now / 15.0)))

        # 电网三相电压
        base_voltage = 3800 + int(25 * math.sin(now / 20.0))
        store.set_value(50, base_voltage)
        store.set_value(51, base_voltage + 8)
        store.set_value(52, base_voltage - 5)

        # PCS 功率与电流/电压
        load_kw = pv_kw * 0.85
        store.set_value(100, max(0, int(load_kw * 10)))
        store.set_value(101, 100 + int(25 * math.cos(now / 25.0)))
        store.set_value(102, 4000 + int(40 * math.sin(now / 28.0)))
        store.set_value(103, 1200 + int(60 * math.sin(now / 30.0)))

        # SOC 模拟充放
        soc = store.get_value(104)
        soc_delta = 2 * math.sin(now / 40.0)
        soc = int(min(950, max(300, soc + soc_delta)))
        store.set_value(104, soc)

        # 温度变化
        warehouse_base += 0.02 * math.sin(now / 60.0)
        store.set_value(93, int(max(200, min(320, warehouse_base))))
        store.set_value(105, store.get_value(93))

        # 堆状态轮换
        phase = int((now / 45.0) % 3)
        storage_state = (phase % 3) + 1  # 1:充电 2:放电 3:待机
        store.set_value(202, storage_state)
        store.set_value(225, 80 + phase * 5)
        store.set_value(226, 90 - phase * 7)

        # 能量累计（0.1 kWh）
        energy_increment = pv_kw * dt / 3600.0 * 10  # 转成 0.1kWh
        daily_energy = (daily_energy + energy_increment) % 100000
        yearly_energy = (yearly_energy + energy_increment) % 600000
        store.set_value(53, int(daily_energy))
        store.set_value(55, int(yearly_energy))

    return updater


class ModbusDataStore:
    """简单的数据存储，支持可选的全局自增。"""

    def __init__(
        self,
        start: int,
        count: int,
        fill: int,
        auto_increment: bool,
        increment_step: int,
        interval: float,
        update_callback: Optional[callable] = None,
    ) -> None:
        self.start = start
        self.count = count
        self.values: List[int] = [fill + i for i in range(count)]
        self.auto_increment = auto_increment
        self.increment_step = increment_step
        self.interval = interval
        self._lock = threading.RLock()
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._update_callback = update_callback

        if auto_increment or update_callback:
            self._thread = threading.Thread(target=self._run, name="ModbusDataStoreAutoIncrement", daemon=True)
            self._thread.start()

    def _run(self) -> None:
        while not self._stop.wait(self.interval):
            with self._lock:
                if self._update_callback:
                    self._update_callback(self)
                elif self.auto_increment:
                    self.values = [(val + self.increment_step) & 0xFFFF for val in self.values]

    def shutdown(self) -> None:
        self._stop.set()
        if self._thread and self._thread.is_alive():
            self._thread.join(timeout=1.0)

    def read(self, start_address: int, quantity: int) -> List[int]:
        with self._lock:
            rel_start = start_address - self.start
            rel_end = rel_start + quantity
            if rel_start < 0 or rel_end > self.count:
                raise ValueError("Illegal data address")
            return self.values[rel_start:rel_end]

    def bump(self, step: int) -> None:
        if not self.auto_increment:
            with self._lock:
                self.values = [(val + step) & 0xFFFF for val in self.values]

    def set_value(self, address: int, value: int) -> None:
        with self._lock:
            rel = address - self.start
            if rel < 0 or rel >= self.count:
                raise ValueError(f"Address {address} out of range [{self.start}, {self.start + self.count - 1}]")
            self.values[rel] = value & 0xFFFF

    def get_value(self, address: int) -> int:
        with self._lock:
            rel = address - self.start
            if rel < 0 or rel >= self.count:
                raise ValueError(f"Address {address} out of range [{self.start}, {self.start + self.count - 1}]")
            return self.values[rel]


class ModbusTCPHandler(socketserver.BaseRequestHandler):
    # FC codes
    READ_HOLDING = 0x03
    READ_INPUT = 0x04

    def handle(self) -> None:
        peer = self.client_address
        logging.info("新连接: %s:%s", peer[0], peer[1])
        try:
            while True:
                # MBAP header is 7 bytes (transaction, protocol, length, unit)
                header = self._recvall(7)
                if not header:
                    break

                transaction_id, protocol_id, length, unit_id = struct.unpack(">HHHB", header)
                if protocol_id != 0:
                    logging.warning("非法协议号: %s (来自 %s)", protocol_id, peer)
                    self._send_exception(transaction_id, unit_id, 0, 1)
                    continue

                pdu = self._recvall(length - 1)
                if not pdu:
                    break

                function_code = pdu[0]
                if function_code in (self.READ_HOLDING, self.READ_INPUT):
                    self._handle_read(transaction_id, unit_id, function_code, pdu[1:])
                else:
                    logging.warning("不支持的功能码: 0x%02X (来自 %s)", function_code, peer)
                    self._send_exception(transaction_id, unit_id, function_code, 0x01)
        except ConnectionResetError:
            logging.info("客户端 %s 断开", peer[0])
        except Exception as exc:
            logging.exception("处理客户端 %s 时异常: %s", peer[0], exc)
        finally:
            logging.info("连接关闭: %s:%s", peer[0], peer[1])

    def _handle_read(self, transaction_id: int, unit_id: int, function_code: int, payload: bytes) -> None:
        if len(payload) != 4:
            self._send_exception(transaction_id, unit_id, function_code, 0x03)
            return

        start_address, quantity = struct.unpack(">HH", payload)
        if quantity == 0 or quantity > 125:
            self._send_exception(transaction_id, unit_id, function_code, 0x03)
            return

        try:
            registers = self.server.data_store.read(start_address, quantity)
            logging.info("读取请求: tx=%s unit=%s func=0x%02X addr=%s qty=%s -> %s",
                         transaction_id, unit_id, function_code, start_address, quantity,
                         registers[:min(6, len(registers))])
        except ValueError:
            logging.warning("非法地址: addr=%s qty=%s", start_address, quantity)
            self._send_exception(transaction_id, unit_id, function_code, 0x02)
            return

        byte_count = len(registers) * 2
        response_pdu = struct.pack(">B", function_code) + struct.pack(">B", byte_count)
        for value in registers:
            response_pdu += struct.pack(">H", value & 0xFFFF)

        self._send_response(transaction_id, unit_id, response_pdu)
        logging.info("已返回: tx=%s 字节数=%s", transaction_id, byte_count)

        # 若未启用自动自增，可在每次请求后手动 bump
        if not self.server.data_store.auto_increment and self.server.bump_step:
            self.server.data_store.bump(self.server.bump_step)

    def _send_response(self, transaction_id: int, unit_id: int, pdu: bytes) -> None:
        length = len(pdu) + 1  # unit id already accounted in header
        mbap = struct.pack(">HHHB", transaction_id, 0, length, unit_id)
        self.request.sendall(mbap + pdu)

    def _send_exception(self, transaction_id: int, unit_id: int, function_code: int, exception_code: int) -> None:
        pdu = struct.pack(">BB", function_code | 0x80, exception_code)
        self._send_response(transaction_id, unit_id, pdu)

    def _recvall(self, size: int) -> bytes:
        data = bytearray()
        while len(data) < size:
            chunk = self.request.recv(size - len(data))
            if not chunk:
                return b""
            data.extend(chunk)
        return bytes(data)


class ModbusTCPServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True

    def __init__(self, server_address, handler_class, data_store, bump_step):
        self.data_store = data_store
        self.bump_step = bump_step
        super().__init__(server_address, handler_class)

    def server_close(self) -> None:
        self.data_store.shutdown()
        super().server_close()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Modbus TCP 测试服务器")
    parser.add_argument("--host", default="0.0.0.0", help="监听地址 (默认: 0.0.0.0)")
    parser.add_argument("--port", type=int, default=1502, help="监听端口 (默认: 1502)")
    parser.add_argument("--start", type=int, default=0, help="寄存器起始地址 (默认: 0)")
    parser.add_argument("--count", type=int, default=400, help="寄存器数量 (默认: 400)")
    parser.add_argument("--fill", type=int, default=0, help="寄存器初始填充值 (默认: 0)")
    parser.add_argument("--auto-increment", action="store_true", help="启用周期自增模拟")
    parser.add_argument("--increment-step", type=int, default=1, help="每次自增步长 (默认: 1)")
    parser.add_argument("--increment-interval", type=float, default=1.0, help="自增间隔秒 (默认: 1.0)")
    parser.add_argument("--bump-step", type=int, default=0, help="若未自动自增，每次读请求后整体加多少 (默认: 0)")
    parser.add_argument("--preset", choices=["zkzf"], help="预置数据集，当前支持: zkzf")
    parser.add_argument(
        "--set",
        action="append",
        default=[],
        metavar="ADDR:VALUE",
        help="手动指定寄存器值，可重复使用。例如 --set 100:123 --set 225:50",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    overrides: Dict[int, int] = {}
    update_callback = None
    if args.preset == "zkzf":
        overrides.update({
            0: 1,  # 运行状态
            35: 1200,  # 光伏功率 (kW*10)
            38: 750,   # DC母线电压
            39: 150,   # DC母线电流
            40: 1800,  # DC母线功率
            50: 3800,  # 电网电压 A
            51: 3812,
            52: 3795,
            53: 5000,  # 当日光伏发电量
            55: 28000, # 年总光伏
            93: 260,   # 仓温 26.0℃
            100: 3450, # PCS总有功功率 (0.1 kW)
            101: 120,  # PCS总无功功率
            102: 4000, # 总电压
            103: 1200, # 总电流
            104: 650,  # 总SOC
            105: 270,  # 平均温度
            107: 345,  # 单体最高电压
            108: 310,  # 单体最低电压
            111: 1,    # 系统状态
            202: 1,    # 堆状态 -> 充电
            225: 80,   # 允许充电功率 (kW)
            226: 90,   # 允许放电功率 (kW)
        })
        update_callback = make_zkzf_update()
        if args.auto_increment:
            logging.info("检测到 --auto-increment，但在 zkzf 预置下改用动态模拟逻辑")
            args.auto_increment = False

    for item in args.set:
        try:
            addr_str, value_str = item.split(":", 1)
            addr = int(addr_str, 0)
            value = int(value_str, 0)
            overrides[addr] = value
        except ValueError:
            raise SystemExit(f"无效的 --set 参数: {item}")

    max_required_addr = args.start + args.count - 1
    if overrides:
        max_override_addr = max(overrides)
        if max_override_addr > max_required_addr:
            needed = max_override_addr - args.start + 1
            logging.info("自动扩展寄存器数量以覆盖地址 %s (调整 count=%s)", max_override_addr, needed)
            args.count = needed

    data_store = ModbusDataStore(
        start=args.start,
        count=args.count,
        fill=args.fill,
        auto_increment=args.auto_increment,
        increment_step=args.increment_step,
        interval=args.increment_interval,
        update_callback=update_callback,
    )

    for addr, value in overrides.items():
        try:
            data_store.set_value(addr, value)
        except ValueError as exc:
            logging.error("设置寄存器值失败: %s", exc)
            raise SystemExit(1)

    server = ModbusTCPServer((args.host, args.port), ModbusTCPHandler, data_store, args.bump_step)

    logging.info(
        "Modbus TCP 测试服务器启动: %s:%s (寄存器范围 %s ~ %s)",
        args.host,
        args.port,
        args.start,
        args.start + args.count - 1,
    )
    logging.info("功能码: 0x03 / 0x04；按 Ctrl+C 停止")

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        logging.info("收到退出信号，正在关闭...")
    finally:
        server.shutdown()
        server.server_close()
        logging.info("服务器已停止")


if __name__ == "__main__":
    main()
